import argparse
import yaml
import torch
from trainer import Trainer
from validator import Validator
from model import create_model
from utils import load_dataset


def load_config(config_path):
	"""加载配置文件"""
	with open(config_path, 'r', encoding='utf-8') as f:
		config = yaml.safe_load(f)
	return config


def main():
	parser = argparse.ArgumentParser(description='古诗生成模型训练与验证')
	parser.add_argument('--config', type=str, default='config.yaml', help='配置文件路径')
	parser.add_argument('--mode', type=str, choices=['train', 'validate'], required=True,
						help='运行模式：train 或 validate')
	parser.add_argument('--model_path', type=str, default=None, help='验证时使用的模型路径（仅validate模式有效）')
	args = parser.parse_args()

	# 加载配置
	config = load_config(args.config)

	if args.mode == 'train':
		# 训练模式
		trainer = Trainer(config)
		trainer.train()

	elif args.mode == 'validate':
		# 验证/生成模式
		validator = Validator(config, model_path=args.model_path)
		validator.validate()


if __name__ == "__main__":
	main()